#include <windows.h>

#include <ace.h>
#include <game.h>
#include <core.h>
#include <utils.h>
#include <msg.h>
#include <pe.h>

#define NTDLL_DYNAMIC_LINK_IMPL
#include <ntdll.h>

#include <main.h>

HMODULE this_module;
size_t unload_ctr = 0;

void unload_ctr_inc() {
    unload_ctr++;
}

void unload_ctr_dec() {
    unload_ctr--;
    if (unload_ctr == 0) {
        void *p_free_library = GetProcAddress(GetModuleHandleA("kernel32.dll"), "FreeLibrary");
        CreateThread(NULL, 0, p_free_library, this_module, 0, NULL);
    }
}

BOOL WINAPI DllMain(HINSTANCE instance, DWORD reason, LPVOID reserved) {
    this_module = instance;
    return TRUE;
}

void *run() {
    void *entry_point = NULL;

    // Dynamically link functions from ntdll
    ntdll_link();

    // Detect which game the user is trying to run
    struct game_data game;
    game_detect(&game);
    
    // Handle the static component (shell)
    switch (game.shell_sub) {
    case SHELL_TX:
        // Unpack shell data
        HMODULE shell_module = ace_load_shell_module();

        size_t shell_data_size;
        void *shell_data = ace_unpack_shell_data(game.shell_args.tx.txs_section_name, shell_module, &shell_data_size);

        // Setup TX
        core_setup_tx(&game, shell_data, shell_data_size);

        // No break here 
    case SHELL_SKIP:
        // No nice entrypoint, we'll have to make one ourselves
        entry_point = game_generate_unity_loader();
        break;

    case SHELL_HAYSTACK:
        // Load shell module anyway, won't be used
        ace_load_shell_module();

        // Unpack the game executable
        LARGE_INTEGER frequency, unpack_start, unpack_end;
        QueryPerformanceFrequency(&frequency);
        QueryPerformanceCounter(&unpack_start);

        ace_unpack_sections();

        QueryPerformanceCounter(&unpack_end);
        if (utils_env_enabled("JADEITE_UNPACK_TIME")) {
            uint64_t elapsed = unpack_end.QuadPart - unpack_start.QuadPart;
            double lfElapsed = (double)elapsed / (double)frequency.QuadPart;
            msg_info_a("Unpacked in %.3lf seconds (%llu QPC ticks)", lfElapsed, elapsed);
        }

        // Find and unpack exe data
        size_t exe_data_size;
        void *exe_data = ace_unpack_exe_data(&exe_data_size);

        // Run Haystack
        core_haystack(exe_data, exe_data_size);

        // Run game-specific entrypoint search
        entry_point = game.shell_args.haystack.get_entry_point();
        break;
    }

    // Handle the runtime component (base)
    switch (game.base_sub) {
    case BASE_FULL:
        // Hook the required api
        const char *symbol;
        void *target;

        switch (game.base_args.full.api) {
        case BASE_API_3:
            symbol = "InitAceClient3";
            target = &core_ace_init3;
            break;
        }

        ace_load_runtime_module(game.base_args.full.base_name, symbol, target);

        // Prevent unloading
        unload_ctr_inc();

        break;

    case BASE_DRV:
        // Hook the single api in drv
        ace_load_runtime_module(game.base_args.drv.drv_name, "CreateObject", &core_drv_co);

        // Prevent unloading
        unload_ctr_inc();

        break;

    case BASE_NONE:
        // none
        break;
    }

    return entry_point;
}
